package com.yuer.image;

import us.codecraft.webmagic.ResultItems;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.pipeline.Pipeline;

import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.Iterator;
import java.util.Map;

/**
 * 获取章节
 */
public class ImagePipeline implements Pipeline {

    private String rootPath;

    public ImagePipeline(String rootPath) {
        this.rootPath = rootPath;
    }

    @Override
    public void process(ResultItems resultItems, Task task) {
        String sourceUrl = resultItems.getRequest().getUrl();
        System.out.println("正在抓取...: " + sourceUrl);

        Iterator iterator = resultItems.getAll().entrySet().iterator();

        while(iterator.hasNext()) {
            Map.Entry<String, String> entry = (Map.Entry)iterator.next();

            String fileName = entry.getKey();
            String fileUrl = entry.getValue();

            URL url;
            InputStream inputStream = null;
            OutputStream outputStream = null;

            try {
                url = new URL(fileUrl);

                HttpURLConnection connection = (HttpURLConnection) url.openConnection();
                connection.setRequestProperty("User-Agent", "Mozilla/5.0 (Windows NT 6.1; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.19 Safari/537.36");
                connection.setConnectTimeout(5 * 1000);

                inputStream = connection.getInputStream();

                outputStream = new FileOutputStream(rootPath + fileName);

                byte[] buffer = new byte[1024];
                int length;

                while ((length = inputStream.read(buffer)) != -1) {
                    outputStream.write(buffer, 0, length);
                }
                inputStream.close();
                outputStream.close();

            } catch (Exception e) {
                System.out.println(e);
            } finally {
                try {
                    if (inputStream != null) {
                        inputStream.close();
                    }
                    if (outputStream != null) {
                        outputStream.close();
                    }
                } catch (Exception e) {
                    System.out.println(e);
                }
            }
        }
    }
}
